module ModuleSteady

    use Globals
    use ModuleInitialize
    use ModuleSolve_GLO_LOC
    use ModuleAggregates
    use ModuleSaving

    contains

	!-----------------------------------------------------------------------!
	! COMMENTS
	!-----------------------------------------------------------------------!
	! E_eval: function that returns steady-state errors at x0
	!         input is NQ-by-1 and output also NQ-by-1
	!
	! JACmat is a NQ-by-NQ matrix -> stores initial Jacobian
	!
	! NTOT is total number of unknows 0> in this case, NTOT = NQ
    !
    ! See Miranda and Fackler Chapter 3.4

	!-----------------------------------------------------------------------!
	!  							MODEL										!
	!-----------------------------------------------------------------------!
	subroutine ComputeSteady
	
	implicit none

    ! Local variable declarations
    real(8), dimension(NQ,NQ) :: Bjinvk, Bjinvkp1
    real(8), dimension(NQ) :: x_in
    real(8) :: Jdstep
    real(8) :: xk(NQ), xkp1(NQ), fk(NQ), fkp1(NQ), Bk(NQ,NQ), Btf(NQ), dx(NQ), fik(NQ), xik(NQ)
    real(8) :: Bkp1(NQ,NQ)
    real(8) :: wwx
    integer :: ix, it, itbs
    real(8) :: eax(1), eak, eikm, ek, eik, ekp1, errf(1), error_old
	
    ! error at initial point x0
	ED0 = ED_eval(x0)		
	
    !-----------------------------------------------------------------!
	! Jacobian
	Jdstep = 1e-5; JACmat = 0.0D0;					! step size; initialize Jacobian
	do ix = 1, NQ									! loop over the equilibrium objects
		x_in     = x0								! initial guess
		x_in(ix) = x0(ix)*(1.0D0+Jdstep)			! increase marginally initial guess of one eqm object at a time
		EDx      = ED_eval(x_in)					! evaluate error after change

		JACmat(:,ix) = (EDx-ED0)/(Jdstep*x0(ix))	! fill Jacobian
	enddo
	!-----------------------------------------------------------------!
	
	!-----------------------------------------------------------------!
	! Iteration
	wwx = 0.50D0    						! updating speed

    index_exit = 0; error_old = 100D0;		! initializations
	
	xk     = x0								! initial guess
	Bjinvk = inv_FET(JACmat)				! inverse of Jacobian	
	fk     = ED0							! error at initial guess
	eax    = maxval(abs(fk));				! max absolute error at initial guess
	ek     = eax(1)							! max absolute error at initial guess
		
	do it = 1, maxitq
	
		!!! new guess
		call matvec_FET(Bjinvk,fk,Btf)		! apply updating rule; matvec_FET(A,V,Y)
		xkp1 = xk - Btf						! apply updating rule
		xkp1 = wwx*xk + (1.0D0-wwx)*xkp1	! slow updating
		
		dx   = xkp1 - xk					! difference between new and previous guess
		
		eikm = 99900000000.0D0				! initialize error from previous backstep
		
		!!! backstep check
		do itbs = 1, maxitqbs

			xik   = xk + dx						! new guess
            xik(2) = max(xik(2),1e-8)			! bound second element of guess: interest rate positive
            xik(1) = min(max(xik(1),lbdlb),1D0)	! bound first element of guess: tax level 

			fik   = ED_eval(xik)				! evaluate at new guess			
			
			eax = maxval(abs(fik))				! max absolute error at new guess
			eik = eax(1)						! max absolute error at new guess
			
			if ( eik < ek ) then 				! this was a good step, move forward
				exit
			elseif ( eikm < eik ) then			! last step was better, undo last shrinking step and exit
				dx  = 2.0D0*dx								
				exit
			endif
			
			! if you reached here, (eik > ek) but (eik < eikm) -> keep shrinking
			eikm = eik							! update error from backstepping
			dx   = dx/2.0D0 					! shrink		
				
		enddo  ! end loop in itbs

		xik = xk + dx							! updated guess
        xik(2) = max(xik(2),1e-8)				! bound second element of guess: interest rate positive
        xik(1) = min(max(xik(1),lbdlb),0.8d0)	! bound first element of guess: tax level
		
		xkp1 = xik; fkp1 = fik; ekp1 = eik;		! update guess, error, max absolute error
		
		errf   = maxval(abs(fkp1))				! max absolute error
		print *, '****************************'
        print *, '** lambda, r = ', xik
		print *, '**  it = ', it, ', errf = ', fik
		print *, 'itbs = ', itbs
		print *, '****************************'

        if ((errf(1)>error_old) .or. (errmu>1e-8)) index_exit = index_exit + 1	! add to index_exit if step was no improvement or measure has not converged
        error_old = errf(1)														! update old error
        
        if (errf(1) < tolf) then					! error sufficiently small
			print *, 'equilibrium found'
			exit
        elseif (index_exit==max_exit) then			! maximum number of unsuccessful steps reached
            print *, 'cant find eqm'
            exit
        else
			!!! update x and B			
			Bjinvkp1 = Update_JACOBIAN(xk,xkp1,fk,fkp1,Bjinvk)			
			
			!!! rename stuff
			xk     = xkp1
			fk     = fkp1
			Bjinvk = Bjinvkp1
			ek     = ekp1
		endif	
    	
	enddo  ! end loop it

	index_it=it

    end subroutine ComputeSteady

	!-----------------------------------------------------------------!
	
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!! 				  Update JACOBIAN				!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
	function Update_JACOBIAN(xk,xkp1,fk,fkp1,Bk)  result(Bkp1)
	
	implicit none
    
    ! Local variable delarations
	real(8), intent(in) :: xk(NQ), xkp1(NQ), fk(NQ), fkp1(NQ), Bk(NQ,NQ)
	real(8)             :: Bkp1(NQ,NQ)
	
	real(8) :: dfk(NQ), ukvec(NQ), dkvec(NQ), Btf(NQ), Bnum1(NQ,1), Bnum2(1,NQ), Bnum(NQ,NQ), dkvecT(1,NQ), denB
	
	!!! update x and B
	dfk   = fkp1 - fk					! change in error
	call matvec_FET(Bk,dfk,ukvec)		! see Miranda Fackler page 38
	dkvec = xkp1 - xk					! change in guess

	!!! numerator computations
	Bnum1(:,1)  = dkvec - ukvec  ! Bnum1 = dk-uk -> a NQ1 "matrix"

	dkvecT(1,:) = dkvec						! transpose		
	call matmul_FET(dkvecT,Bk,Bnum2) 	    ! this is Bnum2 = dk^T * Bk  -> a 1xN "matrix"   ~~ matmul_FET(A,B,C)			
	call matmul_FET(Bnum1,Bnum2,Bnum)       ! this is Bnum = (dk-uk)*(dk^T*Bk) -> a NQN matrix

	!!! denominator computations
	denB  = sum(dkvec*ukvec)

	!!! update Binv
	Bkp1 = Bk + (Bnum/denB)				
	
	end function Update_JACOBIAN


    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!!!                  Compute errors                !!!!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    function ED_eval(xx)  result(ED01)

    implicit none
	
    ! Local variable declarations
    real(8), intent(in) :: xx(NQ)
    real(8)             :: ED01(NQ)

    !!! Compute the economy
    lambda = xx(1)			! level of tax
    r      = xx(2)			! interest rate

    wge = (1.0D0-alpha)/(r+delta)					! wage
    wge = alpha*(wge**((1D0-alpha)/alpha))          ! wage

    call Solve_GLO_LOC                              ! solve HH problem
    call MeasureStationaryPara                      ! compute stationary measure
    call Aggregation                                ! compute aggregates

    ED01(1) = BC			! error in govt budget constraint
    ED01(2) = asset			! error in asset market clearing

    end function ED_eval

end module ModuleSteady
